#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
自动聊天脚本 - 使用 LangChain 管理对话
监听用户回复并使用大模型自动回复
"""

import asyncio
import os
from typing import Optional
from browser_use import Agent, Browser
from browser_use.llm import ChatAzureOpenAI, ChatOpenAI as BrowserChatOpenAI
from langchain_openai import AzureChatOpenAI, ChatOpenAI
from langchain_core.prompts import ChatPromptTemplate, MessagesPlaceholder
from langchain_core.messages import HumanMessage, AIMessage
from langchain_core.runnables.history import RunnableWithMessageHistory
from langchain_community.chat_message_histories import ChatMessageHistory


class AutoChatAgent:
    """自动聊天代理，使用 LangChain 管理对话"""
    
    def __init__(self):
        self.browser = None
        self.browser_agent = None
        self.conversation_chain = None
        self.memory = None
        self.setup_llm()
    
    def setup_llm(self):
        """初始化 LangChain 对话链"""
        print("正在配置 LangChain 对话管理...")
        

        # self.llm = AzureChatOpenAI(
        #     model="gpt-4o",
        #     api_key=os.getenv("PEACH_AZURE_OPENAI_API_KEY"),
        #     azure_endpoint=os.getenv("PEACH_AZURE_OPENAI_ENDPOINT"),
        #     azure_deployment="gpt-4o",
        #     api_version="2024-08-01-preview",
        #     temperature=0.7,
        # )

        # 创建阿里云 Qwen LLM for LangChain
        api_key = os.getenv("ALIBABA_CLOUD")
        base_url = 'https://dashscope.aliyuncs.com/compatible-mode/v1'
        self.llm = ChatOpenAI(
            model='qwen-plus',
            api_key=api_key,
            base_url=base_url,
            temperature=0.7,
        )
        
        # 创建对话记忆
        self.memory = ChatMessageHistory()
        
        # 创建对话提示模板
        self.prompt = ChatPromptTemplate.from_messages([
            ("system", "你是一个CTO，目前团队管理又一些困难，你正在咨询一个管理专家，你需要他帮你解决团队问题，你需要回他管理专家的问题，每次只回复一句话"),
            MessagesPlaceholder(variable_name="history"),
            ("human", "{input}"),
        ])
        
        # 创建对话链
        self.conversation_chain = self.prompt | self.llm
        
        print("LangChain 对话链已初始化")
    
    async def setup_browser(self):
        """初始化浏览器和 Browser Use Agent"""
        print("\n正在配置 Browser Use Agent...")
        
        # 创建 Browser Use 的 LLM
        browser_llm = ChatAzureOpenAI(
            model="gpt-4o",
            api_key=os.getenv("PEACH_AZURE_OPENAI_API_KEY"),
            azure_endpoint=os.getenv("PEACH_AZURE_OPENAI_ENDPOINT"),
            azure_deployment="gpt-4o",
        )
        # api_key = os.getenv("ALIBABA_CLOUD")
        # base_url = 'https://dashscope.aliyuncs.com/compatible-mode/v1'
        # browser_llm = BrowserChatOpenAI(model='qwen-plus', api_key=api_key, base_url=base_url)
        
        # 创建 Browser 实例
        self.browser = Browser(
            headless=False,
            window_size={'width': 1000, 'height': 600},
            keep_alive=True,
        )
        
        print("Browser 已创建")
        return browser_llm
    
    async def send_initial_message(self, browser_llm):
        """发送初始消息"""
        print("\n=== 发送初始消息 ===")
        
        # 初始消息内容
        initial_message = "我是CTO,目前管理60人，我现在太忙了，需要你帮我解决团队问题"
        
        task = f"""
            请执行以下任务:
            
            1. 打开网址: http://134.175.215.235 ，等待10秒, 等待页面加载完成。
            2. 使用用户名 admin 和密码 admin123 登录，等待登录完成。
            3. 点击"新会话"按钮
            4. 在聊天输入框中输入消息"{initial_message}"
            5. 发送消息
            """
        
        agent = Agent(
            task=task,
            llm=browser_llm,
            browser=self.browser,
            use_vision=True,
            max_actions_per_step=5,
            max_failures=3,
        )
        
        history = await agent.run(max_steps=15)
        
        if history.is_successful():
            print("初始消息发送成功")
            # 将初始消息保存到对话历史（作为AI的消息，因为是我们的CTO角色发送的）
            self.memory.add_message(AIMessage(content=initial_message))
            print(f"已保存初始消息到对话历史: {initial_message}")
            return True
        else:
            print("初始消息发送失败")
            return False
    
    async def wait_and_extract_reply(self, browser_llm) -> Optional[str]:
        """等待并提取AI助手的回复"""
        print("\n=== 等待AI助手回复 ===")
        
        task = """
            请执行以下任务:
            
            1. 等待30秒，让AI助手有时间生成回复
            2. 在聊天界面中找到AI助手的最新回复消息（不是我们自己发送的消息）
            3. 提取AI助手的回复内容并返回
            
            重要提示：
            - 聊天界面中有两种消息：我们发送的问题 和 AI助手的回复
            - 只提取AI助手的回复，不要提取我们自己发送的消息
            - AI助手的回复通常比较长，包含建议和分析
            """
        
        agent = Agent(
            task=task,
            llm=browser_llm,
            browser=self.browser,
            use_vision=True,
            max_actions_per_step=5,
            max_failures=3,
        )
        
        history = await agent.run(max_steps=10)
        
        if history.is_successful():
            result = history.final_result()
            if result:
                print(f"提取到AI助手回复: {result}")
                return result
        
        print("未能提取到AI助手回复")
        return None
    
    async def send_reply(self, browser_llm, message: str):
        """发送回复消息"""
        print(f"\n=== 发送回复: {message} ===")
        
        task = f"""
            请执行以下任务:
            
            1. 在聊天输入框中输入消息: {message}
            2. 发送消息
            """
        
        agent = Agent(
            task=task,
            llm=browser_llm,
            browser=self.browser,
            use_vision=True,
            max_actions_per_step=5,
            max_failures=3,
        )
        
        history = await agent.run(max_steps=10)
        
        if history.is_successful():
            print("回复发送成功")
            return True
        else:
            print("回复发送失败")
            return False
    
    def generate_response(self, user_message: str) -> str:
        """使用 LangChain 生成回复"""
        print(f"\n使用 LangChain 生成回复...")
        print(f"收到管理专家的消息: {user_message[:100]}..." if len(user_message) > 100 else f"📥 收到管理专家的消息: {user_message}")
        
        # 先获取当前的历史记录（不包含当前消息）
        current_history = self.memory.messages.copy()
        
        # 使用对话链生成回复（传入历史记录，不包含当前消息）
        response = self.conversation_chain.invoke({
            "input": user_message,
            "history": current_history
        })
        
        # 提取回复内容
        response_text = response.content
        
        # 现在将用户消息和AI回复都添加到历史（按正确顺序）
        self.memory.add_message(HumanMessage(content=user_message))
        self.memory.add_message(AIMessage(content=response_text))
        
        print(f"生成的回复: {response_text}")
        print(f"已保存本轮对话到历史（管理专家消息 + CTO回复）")
        return response_text
    
    async def auto_chat_loop(self, max_rounds: int = 5):
        """自动聊天循环"""
        print(f"\n{'='*60}")
        print(f"开始自动聊天，最多进行 {max_rounds} 轮对话")
        print(f"{'='*60}")
        
        browser_llm = await self.setup_browser()
        
        # 发送初始消息
        if not await self.send_initial_message(browser_llm):
            print("初始化失败，退出")
            return
        
        # 进行多轮对话
        for round_num in range(1, max_rounds + 1):
            print(f"\n{'='*60}")
            print(f"第 {round_num} 轮对话")
            print(f"{'='*60}")
            
            # 等待并提取AI助手回复
            ai_assistant_reply = await self.wait_and_extract_reply(browser_llm)
            
            if not ai_assistant_reply:
                print("未收到AI助手回复，等待下一轮...")
                await asyncio.sleep(5)
                continue
            
            # 使用 LangChain 生成回复（基于AI助手的回复）
            our_response = self.generate_response(ai_assistant_reply)
            
            # 发送回复
            if await self.send_reply(browser_llm, our_response):
                print(f"第 {round_num} 轮对话完成")
            else:
                print(f"第 {round_num} 轮发送失败")
                break
            
            # 短暂等待，避免过快
            await asyncio.sleep(3)
        
        print(f"\n{'='*60}")
        print("自动聊天结束")
        print(f"{'='*60}")
        
        # 显示对话历史
        print("\n=== 对话历史 ===")
        for i, msg in enumerate(self.memory.messages, 1):
            role = "管理专家(网页AI)" if isinstance(msg, HumanMessage) else "CTO(我们的AI)"
            print(f"{i}. {role}: {msg.content[:100]}..." if len(msg.content) > 100 else f"{i}. {role}: {msg.content}")
    
    async def run(self, max_rounds: int = 5):
        """运行自动聊天"""
        try:
            await self.auto_chat_loop(max_rounds)
            
            # 保持浏览器打开
            print("\n浏览器将保持打开 60 秒，让你查看结果...")
            print("按 Ctrl+C 可以立即关闭")
            try:
                await asyncio.sleep(60)
            except KeyboardInterrupt:
                print("\n收到中断信号...")
        
        except Exception as e:
            print(f"\n错误: {e}")
            import traceback
            traceback.print_exc()
            
            print("\n浏览器将保持打开 60 秒...")
            try:
                await asyncio.sleep(60)
            except KeyboardInterrupt:
                pass
        
        finally:
            if self.browser:
                try:
                    await self.browser.stop()
                    print("\n浏览器已关闭")
                except:
                    pass


async def main():
    """主函数"""
    # 创建自动聊天代理
    agent = AutoChatAgent()
    
    # 运行自动聊天，最多5轮对话
    await agent.run(max_rounds=5)


if __name__ == "__main__":
    try:
        asyncio.run(main())
    except KeyboardInterrupt:
        print("\n程序已退出")
